MODULE sefmm_cheb_utils

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: cheb_get_coeffs,           &
             cheb_get_function,         &
             cheb_truncate_coeffs,      &
             cheb_translate_coeffs,     &
             cheb_trans_trace,          &
             cheb_untrans_trace,        &
             cheb_inner,                &
             cheb_get_nodes,            &
             cheb_get_Tn_ints,          &
             cheb_accurate_order,       &
             cheb_get_1D_coeffs,        &
             cheb_1D_truncate_coeffs

CONTAINS

!------------------------------------------------------------------------------

   FUNCTION cheb_accurate_order(grain,order,expp)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: order
      REAL(REALK),   INTENT(IN)  :: grain,expp
      LOGICAL :: cheb_accurate_order 

      REAL(REALK) :: ratio, s 

      cheb_accurate_order = .TRUE.
      IF (expp < 0.01d0) RETURN

      ratio = expp / 0.01d0
!      s = 10d0 / SQRT(ratio)
      s = 50d0 / SQRT(ratio)

      cheb_accurate_order = .FALSE.
      IF (order == 12) THEN
!         print *, grain, expp, s
         IF ( grain < s ) cheb_accurate_order = .TRUE.
!         cheb_accurate_order = ( grain < s )
      ELSE
         print *, 'order must = 8'; stop
      END IF

   END FUNCTION cheb_accurate_order

!------------------------------------------------------------------------------

   SUBROUTINE cheb_get_function(npts,npts_out,coeffs,func)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts, npts_out
      ! transformed
      REAL(REALK),   INTENT(IN)  :: coeffs(npts,npts,npts)   
      ! untransformed
      REAL(REALK),   INTENT(OUT) :: func(npts_out,npts_out,npts_out)     

      REAL(REALK)   :: Tn(npts_out,npts)
      REAL(REALK)   :: tmp1(npts_out,npts,npts)
      REAL(REALK)   :: tmp2(npts_out,npts_out,npts)
      INTEGER(INTK) :: ix,iy,iz

      CALL get_basis(npts_out,npts,Tn)

!FIXME speed-up contraction ordering!!
      DO iz = 1, npts
         DO iy = 1, npts
            DO ix = 1, npts_out
               tmp1(ix,iy,iz) = DOT_PRODUCT(Tn(ix,:),coeffs(:,iy,iz))
            END DO
         END DO
      END DO

      DO iz = 1, npts
         DO ix = 1, npts_out
            DO iy = 1, npts_out
               tmp2(ix,iy,iz) = DOT_PRODUCT(Tn(iy,:),tmp1(ix,:,iz))
            END DO
         END DO
      END DO

      DO iy = 1, npts_out
         DO ix = 1, npts_out
            DO iz = 1, npts_out
               func(ix,iy,iz) = DOT_PRODUCT(Tn(iz,:),tmp2(ix,iy,:))
            END DO
         END DO
      END DO

   END SUBROUTINE cheb_get_function

!------------------------------------------------------------------------------

   SUBROUTINE cheb_get_coeffs(npts,func,coeffs)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(IN)  :: func(npts,npts,npts)    ! untransformed
      REAL(REALK),   INTENT(OUT) :: coeffs(npts,npts,npts)  ! transformed

      REAL(REALK) :: fac

      SELECT CASE (npts)
         CASE (8)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (16)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (32)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (64)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (128)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (256)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (512)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE (1024)
            CALL cheb_get_fast_coeffs(npts,func,coeffs)
         CASE DEFAULT
            CALL cheb_get_slow_coeffs(npts,func,coeffs)
      END SELECT

      fac = 1d0/(npts*npts*npts)
      coeffs(:,:,:) = coeffs(:,:,:) * 8.0d0 * fac
      coeffs(1,:,:) = coeffs(1,:,:) * half
      coeffs(:,1,:) = coeffs(:,1,:) * half
      coeffs(:,:,1) = coeffs(:,:,1) * half

   END SUBROUTINE cheb_get_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE cheb_truncate_coeffs(npts,coeffs,order)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(IN)  :: coeffs(npts,npts,npts)
      INTEGER(INTK), INTENT(OUT) :: order

      INTEGER(INTK) :: n0,n1, i,j,k
      REAL(REALK)   :: tmp

      n0 = 2
      n1 = 2
      DO WHILE (n1 <= npts)

         tmp = zero
         DO i = n0, n1
            DO j = n0, n1
               DO k = n0, n1
                  tmp = tmp + ABS(coeffs(i,j,k))
               END DO
            END DO
         END DO

         !print *, n0,n1, tmp
         order = n1 
         IF (tmp < 1d-10) RETURN
         n0 = n1+1

         !n1 = n1*2
         n1 = n1+2

      END DO

      order = npts

   END SUBROUTINE cheb_truncate_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE cheb_translate_coeffs(order0,order1,func0,func1)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: order0, order1
      REAL(REALK),   INTENT(IN)  :: func0(order0,order0,order0)
      REAL(REALK),   INTENT(OUT) :: func1(order1,order1,order1)

      REAL(REALK) :: coeffs(order0,order0,order0)
      REAL(REALK) :: trace0, trace1

      CALL cheb_get_coeffs(order0,func0,coeffs)
      CALL cheb_get_function(order0,order1,coeffs,func1)

! checks
!print *, 'translating... check traces!'
!      trace0 = cheb_trans_trace(order0,1d0,coeffs)
!      CALL cheb_get_coeffs(order1,func1,coeffs1)
!      trace1 = cheb_trans_trace(order1,1d0,coeffs1)
!      print *, trace0, trace1

   END SUBROUTINE cheb_translate_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE cheb_1D_truncate_coeffs(npts,coeffs,order)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(IN)  :: coeffs(npts)
      INTEGER(INTK), INTENT(OUT) :: order

      INTEGER(INTK) :: n0,n1, k
      REAL(REALK)   :: tmp

      n0 = 2
      n1 = 2
      DO WHILE (n1 <= npts)

         tmp = zero
         DO k = n0, n1
            tmp = tmp + ABS(coeffs(k))
         END DO

         !print *, n0,n1, tmp
         order = n1 
         IF (tmp < 1d-3) RETURN
         n0 = n1+1

         !n1 = n1*2
         n1 = n1+1

      END DO

      order = npts

   END SUBROUTINE cheb_1D_truncate_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE cheb_get_1D_coeffs(npts,func,coeffs)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(IN)  :: func(npts)    ! untransformed
      REAL(REALK),   INTENT(OUT) :: coeffs(npts)  ! transformed

      REAL(REALK) :: Tn(npts,npts)

      CALL get_basis(npts,npts,Tn)

      coeffs(:) = MATMUL(func(:),Tn)

   END SUBROUTINE cheb_get_1D_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE cheb_get_slow_coeffs(npts,func,coeffs)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(IN)  :: func(npts,npts,npts)    ! untransformed
      REAL(REALK),   INTENT(OUT) :: coeffs(npts,npts,npts)  ! transformed

      REAL(REALK)   :: Tn(npts,npts)
      REAL(REALK)   :: tmp(npts,npts,npts)
      INTEGER(INTK) :: ix,iy,iz

      CALL get_basis(npts,npts,Tn)

      DO iz = 1, npts
         DO iy = 1, npts
            coeffs(:,iy,iz) = MATMUL(func(:,iy,iz),Tn)
         END DO
      END DO

      ! Reorder indices
      DO iz = 1, npts
         DO ix = 1, npts
            tmp(:,ix,iz) = coeffs(ix,:,iz)
         END DO
      END DO

      DO iz = 1, npts
         DO ix = 1, npts
            coeffs(:,ix,iz) = MATMUL(tmp(:,ix,iz),Tn)
         END DO
      END DO

      ! Reorder indices
      DO ix = 1, npts
         DO iy = 1, npts
            tmp(:,iy,ix) = coeffs(iy,ix,:)
         END DO
      END DO

      DO ix = 1, npts
         DO iy = 1, npts
            coeffs(:,iy,ix) = MATMUL(tmp(:,iy,ix),Tn)
         END DO
      END DO

      tmp = coeffs
      ! Reorder indices
      DO iz = 1, npts
         DO iy = 1, npts
            DO ix = 1, npts
               coeffs(ix,iy,iz) = tmp(iz,iy,ix)
            END DO
         END DO
      END DO

   END SUBROUTINE cheb_get_slow_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE cheb_get_fast_coeffs(npts,func,coeffs)

      USE fast_cosine_transform, ONLY: cosft2

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(IN)  :: func(npts,npts,npts)    ! untransformed
      REAL(REALK),   INTENT(OUT) :: coeffs(npts,npts,npts)  ! transformed

      REAL(REALK)   :: tmp1(npts,npts,npts)
      REAL(REALK)   :: tmp2(npts,npts,npts)
      INTEGER(INTK) :: ix,iy,iz

      coeffs = func

      DO iz = 1, npts
         DO iy = 1, npts
            CALL cosft2(npts,1,coeffs(:,iy,iz))
         END DO
      END DO

      ! Reorder indices
      DO iz = 1, npts
         DO ix = 1, npts
            tmp1(:,ix,iz) = coeffs(ix,:,iz)
         END DO
      END DO

      DO iz = 1, npts
         DO ix = 1, npts
            CALL cosft2(npts,1,tmp1(:,ix,iz))
         END DO
      END DO

      ! Reorder indices
      DO ix = 1, npts
         DO iy = 1, npts
            tmp2(:,iy,ix) = tmp1(iy,ix,:)
         END DO
      END DO

      DO ix = 1, npts
         DO iy = 1, npts
            CALL cosft2(npts,1,tmp2(:,iy,ix))
         END DO
      END DO

      ! Reorder indices
      DO iz = 1, npts
         DO iy = 1, npts
            DO ix = 1, npts
               coeffs(ix,iy,iz) = tmp2(iz,iy,ix)
            END DO
         END DO
      END DO

   END SUBROUTINE cheb_get_fast_coeffs

!------------------------------------------------------------------------------

   SUBROUTINE get_basis(npts1,npts2,Tn)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts1, npts2
      REAL(REALK),   INTENT(OUT) :: Tn(0:npts1-1,0:npts2-1)

      INTEGER(INTK) :: n,k

      DO n = 0, npts2-1     ! polynomial
         DO k = 0, npts1-1  ! node
            Tn(k,n) = COS(PI*n*(k+half)/npts1)
         END DO
      END DO

   END SUBROUTINE get_basis

!------------------------------------------------------------------------------

   SUBROUTINE cheb_get_nodes(npts,nodes)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: npts
      REAL(REALK),   INTENT(OUT) :: nodes(0:npts-1)

      INTEGER(INTK) :: k

      DO k = 0, npts-1
         nodes(k) = COS(PI*(k+half)/npts)
      END DO

   END SUBROUTINE cheb_get_nodes

!------------------------------------------------------------------------------

   FUNCTION cheb_inner(npts,lambda,f1,f2)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: npts
      REAL(REALK),   INTENT(IN) :: lambda
      REAL(REALK),   INTENT(IN) :: f1(npts,npts,npts)  ! untransformed
      REAL(REALK),   INTENT(IN) :: f2(npts,npts,npts)  ! untransformed
      REAL(REALK) :: cheb_inner
   
      REAL(REALK)   :: u(npts,npts,npts)
      REAL(REALK)   :: coeffs(npts,npts,npts)
      INTEGER(INTK) :: i,j,k
   
      u = zero
      DO k = 1, npts
         DO j = 1, npts
            DO i = 1, npts
               u(i,j,k) = f1(i,j,k)*f2(i,j,k)
            END DO
         END DO
      END DO

!FIXME we only need 1/8 of the coeffs for the trace !!
      CALL cheb_get_coeffs(npts,u,coeffs)
      cheb_inner = cheb_trans_trace(npts,lambda,coeffs)

!      cheb_inner = cheb_untrans_trace(npts,lambda,u)

   END FUNCTION cheb_inner

!------------------------------------------------------------------------------
!  Compute trace as Gauss-Chebyshev quadrature on 
!  functions discretised on Chebyshev grid

   FUNCTION cheb_untrans_trace(npts,lambda,func)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: npts
      REAL(REALK),   INTENT(IN) :: lambda
      ! Untransformed function
      REAL(REALK),   INTENT(IN) :: func(npts,npts,npts)  
      REAL(REALK) :: cheb_untrans_trace
   
      REAL(REALK)   :: u(npts,npts)
      REAL(REALK)   :: v(npts)
      REAL(REALK)   :: nodes(npts)
      REAL(REALK)   :: wts(npts), trace, x
      INTEGER(INTK) :: i,j,k
   
      CALL cheb_get_nodes(npts,nodes)

      DO i = 1, npts
         wts(i) = (PI/npts)*SQRT(one-nodes(i)*nodes(i))
      END DO

      DO k = 1, npts
         DO j = 1, npts
            u(j,k) = zero
            DO i = 1, npts
               u(j,k) = u(j,k) + wts(i)*func(i,j,k)
            END DO
         END DO
      END DO

      DO k = 1, npts
         v(k) = zero
         DO j = 1, npts
            v(k) = v(k) + wts(j)*u(j,k)
         END DO
      END DO

      trace = zero
      DO k = 1, npts
         trace = trace + wts(k)*v(k)
      END DO

      cheb_untrans_trace = trace*(lambda*half)**3

   END FUNCTION cheb_untrans_trace

!------------------------------------------------------------------------------

   FUNCTION cheb_trans_trace(npts,lambda,cf)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: npts
      REAL(REALK),   INTENT(IN) :: lambda
      ! Transformed function
      REAL(REALK),   INTENT(IN) :: cf(0:npts-1,0:npts-1,0:npts-1)
      REAL(REALK) :: cheb_trans_trace
   
      REAL(REALK)   :: ints(0:npts-1)
      REAL(REALK)   :: trace
      INTEGER(INTK) :: i,j,k
   
      CALL cheb_get_Tn_ints(npts-1,ints)

      trace = zero
      DO i = 0, npts-1, 2
         DO j = 0, npts-1, 2
            DO k = 0, npts-1, 2
               trace = trace + cf(k,j,i)*ints(i)*ints(j)*ints(k)
            END DO
         END DO
      END DO

      cheb_trans_trace = trace*(lambda*half)**3

   END FUNCTION cheb_trans_trace

!------------------------------------------------------------------------------
! Build Tn integrals on [-1,1] 
! \int_{-1}^{1} Tn(x) dx

   SUBROUTINE cheb_get_Tn_ints(order,ints)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: order
      REAL(REALK),   INTENT(OUT) :: ints(0:order)

      INTEGER(INTK) :: i,j,k

      ints(:) = zero
      DO k = 0, order, 2
         ! Integrals of odd Tn are zero
         ints(k) = 2d0*(DBLE(k)/(k**2 -1) - 1d0/(k-1) )
      END DO

   END SUBROUTINE cheb_get_Tn_ints

!------------------------------------------------------------------------------

END MODULE sefmm_cheb_utils
